#include "fun.h"
// #include "file_utils.h"
#include "image_utils.h"
#include <stdio.h>
#include <string.h>
#include <string>

#include "main.h"

static void dump_tensor_attr(rknn_tensor_attr* attr)
{
    std::string shape_str = attr->n_dims < 1 ? "" : std::to_string(attr->dims[0]);
    for (int i = 1; i < attr->n_dims; ++i) {
        shape_str += ", " + std::to_string(attr->dims[i]);
    }

    printf("  index=%d, name=%s, n_dims=%d, dims=[%s], n_elems=%d, size=%d, fmt=%s, type=%s, qnt_type=%s, "
           "zp=%d, scale=%f\n",
           attr->index, attr->name, attr->n_dims, shape_str.c_str(), attr->n_elems, attr->size,get_format_string(attr->fmt),
           get_type_string(attr->type), get_qnt_type_string(attr->qnt_type), attr->zp, attr->scale);
}



int init_model(model_context* model_ctx, const char* model_path)
{
    int ret;
    rknn_context ctx = 0;

    //load .rknn
    ret = rknn_init(&ctx, (char*)model_path, 0, 0, NULL);
    if (ret<0)
    {
        printf("rknn init failure: %s\n",model_path);
    }

    // get IO num
    rknn_input_output_num io_num;
    ret = rknn_query(ctx, RKNN_QUERY_IN_OUT_NUM, &io_num, sizeof(io_num));
    if (ret != RKNN_SUCC)
    {
        printf("rknn_query(IO_NUM) fail! ret=%d\n", ret);
        return -1;
    }
    printf("model input num: %d, output num: %d\n", io_num.n_input, io_num.n_output);

    // input info
    rknn_tensor_attr input_attrs[io_num.n_input];
    memset(input_attrs, 0, sizeof(input_attrs));
    for (int i = 0; i < io_num.n_input; i++)
    {
        input_attrs[i].index = i;
        ret = rknn_query(ctx, RKNN_QUERY_INPUT_ATTR, &(input_attrs[i]), sizeof(rknn_tensor_attr));
        if (ret != RKNN_SUCC)
        {
            printf("rknn_query fail! ret=%d\n", ret);
            return -1;
        }
        dump_tensor_attr(&(input_attrs[i]));
    }

    // output info
    rknn_tensor_attr output_attrs[io_num.n_output];
    memset(output_attrs, 0, sizeof(output_attrs));
    for (int i = 0; i < io_num.n_output; i++)
    {
        output_attrs[i].index = i;
        ret = rknn_query(ctx, RKNN_QUERY_OUTPUT_ATTR, &(output_attrs[i]), sizeof(rknn_tensor_attr));
        if (ret != RKNN_SUCC)
        {
            printf("rknn_query fail! ret=%d\n", ret);
            return -1;
        }
        dump_tensor_attr(&(output_attrs[i]));
    }

    // set to context
    model_ctx->rknn_ctx = ctx;
    
    model_ctx->io_num = io_num;

    model_ctx->input_attrs = (rknn_tensor_attr *)malloc(io_num.n_input * sizeof(rknn_tensor_attr));
    memcpy(model_ctx->input_attrs, input_attrs, io_num.n_input * sizeof(rknn_tensor_attr));
    
    model_ctx->output_attrs = (rknn_tensor_attr *)malloc(io_num.n_output * sizeof(rknn_tensor_attr));
    memcpy(model_ctx->output_attrs, output_attrs, io_num.n_output * sizeof(rknn_tensor_attr));

    if (input_attrs[0].fmt == RKNN_TENSOR_NCHW)
    {
        printf("model is NCHW input fmt\n");
        model_ctx->model_channel = input_attrs[0].dims[1];
        model_ctx->model_height = input_attrs[0].dims[2];
        model_ctx->model_width = input_attrs[0].dims[3];
        printf("input image height=%d, input image width=%d, input image channel=%d\n",
                model_ctx->model_height, model_ctx->model_width, model_ctx->model_channel);
    }
    else if (input_attrs[0].fmt == RKNN_TENSOR_NHWC)
    {
        printf("model is NHWC input fmt\n");
        model_ctx->model_height = input_attrs[0].dims[1];
        model_ctx->model_width = input_attrs[0].dims[2];
        model_ctx->model_channel = input_attrs[0].dims[3];
        printf("input image height=%d, input image width=%d, input image channel=%d\n",
                model_ctx->model_height, model_ctx->model_width, model_ctx->model_channel);
    }
    else
    {
        printf("model is UNDEFINED input fmt\n");
        model_ctx->model_height = input_attrs[0].dims[0];
        model_ctx->model_width = input_attrs[0].dims[1];
        printf("input text batch size=%d, input sequence length=%d\n", 
                model_ctx->model_height, model_ctx->model_width);
    }
    printf("\n");
    return 0;
}



int free_model(model_context* model_ctx)
{
    printf("free this model...");
    if (model_ctx->input_attrs != NULL)
    {
        free(model_ctx->input_attrs);
        model_ctx->input_attrs = NULL;
    }
    if (model_ctx->output_attrs != NULL)
    {
        free(model_ctx->output_attrs);
        model_ctx->output_attrs = NULL;
    }
    if (model_ctx->rknn_ctx != 0)
    {
        rknn_destroy(model_ctx->rknn_ctx);
        model_ctx->rknn_ctx = 0;
    }
    printf("finished\n");
    return 0;
}



int load_and_transform_image(image_buffer_t* image_info, const char *image_path)
{
    int ret;
    image_buffer_t original_image;
    
    //read
    memset(&original_image, 0, sizeof(image_buffer_t));
    ret = read_image(image_path, &original_image);
    if(ret < 0 ) return -1;
    // printf("orinal image info: %d  %d  %d  %d \n", 
    //     original_image.width,
    //     original_image.height,
    //     original_image.size,
    //     original_image.format
    // );

    //prepare paras
    memset(image_info, 0, sizeof(image_buffer_t));
    image_info->width = IMG_WH; // well well magic number
    image_info->height = IMG_WH; // man what can i say
    image_info->format = IMAGE_FORMAT_RGB888;
    image_info->size = get_image_size(image_info);
    image_info->virt_addr = (unsigned char*)malloc(image_info->size);
    //resize 900*900 -> 518*518
    ret = convert_image(&original_image, image_info, NULL, NULL, 0);
    if(ret<0) return -1;

    //free original image, save memory
    printf("free original image: ");
    free_image(&original_image);
    return 0;
}



int free_image(image_buffer_t* image_info)
{
    printf("free image...");
    if (image_info->virt_addr != NULL)
    {
        free(image_info->virt_addr);
    }
    printf("done\n");
    return 0;
}


int inference(model_context* model, image_buffer_t* img_in)
{

#ifndef QUANTIZATION
    // #todo
    //extend int8 to float16
    int8_t * old_int8 = (int8_t*)img_in->virt_addr;
    __fp16 * new_fp16 = (__fp16*)malloc(img_in->size*2);
    for(int i=0; i < img_in->size; i++)
    {
        new_fp16[i] = ((__fp16)old_int8[i])/255.0;
    }
#endif

    int ret;
    rknn_input inputs[1];
    rknn_output outputs[1];
    memset(inputs, 0, sizeof(inputs));
    memset(outputs, 0, sizeof(outputs));

    //prepare input
    inputs[0].index = 0;
    inputs[0].type = RKNN_TENSOR_INT8;
    inputs[0].fmt = RKNN_TENSOR_NHWC;
    inputs[0].size = img_in->size;
    inputs[0].buf = img_in->virt_addr;
#ifndef QUANTIZATION
    inputs[0].type = RKNN_TENSOR_FLOAT16;
    inputs[0].size = img_in->size * 2;
    inputs[0].buf = new_fp16;
#endif
    ret = rknn_inputs_set(model->rknn_ctx, 1, inputs);
    if(ret < 0) {
        printf("rknn_input_set fail! ret=%d\n", ret);
        return -1;
    }

    //infer
    ret = rknn_run(model->rknn_ctx, nullptr);
    if (ret < 0) {
        printf("rknn_run fail! ret=%d\n", ret);
        return -1;
    }

    //extract output
    ret = rknn_outputs_get(model->rknn_ctx, 1, outputs, NULL);
    if (ret < 0) {
        printf("rknn_outputs_get fail! ret=%d\n", ret);
        return -1;
    }

#ifndef QUANTIZATION
    // #todo
    //extend int8 to float16
    __fp16 * old_fp16 = (__fp16*)outputs[0].buf;
    int8_t * new_int8 = (int8_t*)malloc(outputs[0].size/2+1);
    // printf("output size is %d\n",outputs[0].size);
    for(int i=0; i < outputs[0].size/2; i++)
    {
        new_int8[i] = (int8_t)(old_fp16[i]*255.0);
    }
#endif

    image_buffer_t img;
    img.width = IMG_WH;
    img.height = IMG_WH;
    img.width_stride = 0;
    img.height_stride = 0;
    img.format = IMAGE_FORMAT_RGB888;
    img.virt_addr = (unsigned char*)outputs[0].buf;
#ifndef QUANTIZATION
    img.virt_addr = (uint8_t*)new_int8;
#endif
    img.size = 3*IMG_WH*IMG_WH;
    


    write_image("/home/orangepi/Install_AnomalyCLIP/out.png", &img);
    printf("free temporary type-convert buffer\n");

#ifndef QUANTIZATION
    free(new_fp16); free(new_int8);
#endif
    // free_image(img);
    rknn_outputs_release(model->rknn_ctx, 1, outputs); //#todo
    printf("\n~~~~~what can i say, mamba out~~~~~\n");
    return 0;
}



int inference_advanced(model_context* model, image_buffer_t* input)
{
    using namespace std;
    int ret;

    const auto num_inputs = model->io_num.n_input;
    rknn_tensor_mem *input_mems[model->io_num.n_input];

    for(size_t i=0; i<num_inputs; i++)
    {
        // model->input_attrs[i].type = RKNN_TENSOR_UINT8; //wtf
        input_mems[i] = rknn_create_mem(
            model->rknn_ctx,
            model->input_attrs[i].size_with_stride
        );
        // stride is 520 but width is 518, so we need to be careful
        auto height = model->input_attrs[i].dims[1]; //NHWC
        auto width = model->input_attrs[i].dims[2];
        auto channel = model->input_attrs[i].dims[3];
        auto stride = model->input_attrs[i].w_stride;
        
        // pointer
        uint8_t* src_ptr = input->virt_addr;
        uint8_t* dst_ptr = (uint8_t*)input_mems[i]->virt_addr;
        
        // copy
        auto src_wc_elems = width*channel*sizeof(uint8_t);
        auto dst_wc_elems = stride*channel*sizeof(uint8_t);
        
        for (int b = 0; b < model->input_attrs[i].dims[0]; b++) {
            for (int h = 0; h < height; ++h) {
                memcpy(dst_ptr, src_ptr, src_wc_elems);
                src_ptr += src_wc_elems;
                dst_ptr += dst_wc_elems;
            }
        }
        ret = rknn_set_io_mem(
            model->rknn_ctx,
            input_mems[i],
            &(model->input_attrs[i])
        );
        if(ret<0){
            printf("rknn_set_io_mem failed, ret=%d\n", ret);
            return -1;
        }

    }

    const auto num_outputs = model->io_num.n_output;
    rknn_tensor_mem *output_mems[num_outputs];

    for(size_t i=0; i<num_outputs; i++)
    {
        output_mems[i] = rknn_create_mem(
            model->rknn_ctx,
            model->output_attrs[i].size_with_stride
        );
        ret = rknn_set_io_mem(
            model->rknn_ctx,
            output_mems[i],
            &(model->output_attrs[i])
        );
        if(ret < 0){
            printf("rknn_set_io_mem fail! ret=%d\n", ret);
            return -1;
        }
    }
    return -1919810;
}
